# coding:utf-8
# Author : hiicy redldw
# Date : 2019/04/04
import os

from torch import nn

"""
######## 初始权重化 ################
class discriminator(nn.Module):

    def __init__(self, dataset = 'mnist'):
        super(discriminator, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(self.input_dim, 64, 4, 2, 1),
            nn.ReLU(),
        )
        self.fc = nn.Sequential(
            nn.Linear(32, 64 * (self.input_height // 2) * (self.input_width // 2)),
            nn.BatchNorm1d(64 * (self.input_height // 2) * (self.input_width // 2)),
            nn.ReLU(),
        )
        self.deconv = nn.Sequential(
            nn.ConvTranspose2d(64, self.output_dim, 4, 2, 1),
            #nn.Sigmoid(),         # EBGAN does not work well when using Sigmoid().
        )
        initialize_weights(self)

    def forward(self, input):
        pass

def initialize_weights(net):
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()
        elif isinstance(m, nn.ConvTranspose2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            m.weight.data.normal_(0, 0.02)
        m.bias.data.zero_()


def init_weights(m):
    print(m)
    if type(m) == nn.Linear:
        m.weight.data.fill_(1.0)
        print(m.weight)
net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)
"""

########## pytorch 加载(.pth)格式的模型 ###########
import torch
import torchvision.models as models
pthfile=r"D:\hiicy\Google\resnet18-5c106cde.pth"
# net = torch.load(pthfile)#这里只保存了参数
net = models.resnet18(pretrained=False)
net.load_state_dict(torch.load(pthfile))
print(net)
####### 分层学习率 ################

####### 打印模型尺寸 ############
def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p")/1e6)
    os.remove('temp.p')